package com.strp.mptcp.path;

import com.strp.mptcp.path.Subflow;
import io.netty.bootstrap.Bootstrap;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.ssl.SslContext;
import lombok.Getter;

import java.net.SocketAddress;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * PathManager：管理所有子流的生命周期、状态和重连
 * @author 曹行阳
 * @date 2025/6/12 11:31
 * @version 0.1
 */
public class PathManager {
    private final EventLoopGroup group;
    private final SslContext sslContext;
    private final ConcurrentHashMap<Integer, Subflow> subflows = new ConcurrentHashMap<>();
    private final AtomicInteger nextPathId = new AtomicInteger(0);
    @Getter
    private final List<Subflow> healthySubflows = new CopyOnWriteArrayList<>();

    // 心跳相关配置
    private static final long HEARTBEAT_INTERVAL_MS = 5000;
    private static final long HEARTBEAT_TIMEOUT_MS = 3000;
    private static final int MAX_RECONNECT_ATTEMPTS = 3;
    private static final long RECONNECT_DELAY_MS = 1000;

    // Subflow 配置参数
    private static final int MAX_RECEIVE_QUEUE_SIZE = 1000;
    private static final int READ_TIMEOUT = 30;
    private static final int WRITE_TIMEOUT = 30;
    private static final int CONNECTION_TIMEOUT = 10;
    private static final int MAX_RETRIES = 3;

    private final ScheduledExecutorService heartbeatExecutor;
    private final ConcurrentHashMap<Integer, AtomicLong> lastHeartbeatTime = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<Integer, Integer> reconnectAttempts = new ConcurrentHashMap<>();

    public PathManager(EventLoopGroup group, SslContext sslContext) {
        this.group = group;
        this.sslContext = sslContext;
        this.heartbeatExecutor = Executors.newSingleThreadScheduledExecutor();
        startHeartbeatCheck();
    }

    private void startHeartbeatCheck() {
        heartbeatExecutor.scheduleAtFixedRate(() -> {
            long currentTime = System.currentTimeMillis();
            for (Subflow subflow : subflows.values()) {
                AtomicLong lastTime = lastHeartbeatTime.get(subflow.getPathId());
                if (lastTime != null && currentTime - lastTime.get() > HEARTBEAT_TIMEOUT_MS) {
                    handlePathFailure(subflow.getPathId());
                }
            }
        }, HEARTBEAT_INTERVAL_MS, HEARTBEAT_INTERVAL_MS, TimeUnit.MILLISECONDS);
    }

    public void createAndConnectPaths(List<SocketAddress> addresses) {
        Bootstrap b = new Bootstrap()
                .group(group)
                .channel(NioSocketChannel.class)
                .option(io.netty.channel.ChannelOption.SO_KEEPALIVE, true)
                .option(io.netty.channel.ChannelOption.TCP_NODELAY, true);

        for (SocketAddress address : addresses) {
            int pathId = nextPathId.getAndIncrement();
            Subflow subflow = new Subflow(
                    pathId,
                    address,
                    b,
                    sslContext,
                    MAX_RECEIVE_QUEUE_SIZE,
                    READ_TIMEOUT,
                    WRITE_TIMEOUT,
                    CONNECTION_TIMEOUT,
                    MAX_RETRIES
            );
            subflows.put(pathId, subflow);
            connectSubflow(subflow);
        }
    }

    private void connectSubflow(Subflow subflow) {
        subflow.connect();
        // 由于 connect() 方法现在是异步的，我们需要等待连接完成
        group.schedule(() -> {
            if (subflow.isActive()) {
                onPathConnected(subflow);
            } else {
                handlePathFailure(subflow.getPathId());
            }
        }, CONNECTION_TIMEOUT, TimeUnit.SECONDS);
    }

    private void onPathConnected(Subflow subflow) {
        healthySubflows.add(subflow);
        lastHeartbeatTime.put(subflow.getPathId(), new AtomicLong(System.currentTimeMillis()));
        reconnectAttempts.put(subflow.getPathId(), 0);
        subflow.startHeartbeat();
    }

    public void onPathFailure(int pathId) {
        handlePathFailure(pathId);
    }

    private void handlePathFailure(int pathId) {
        Subflow failedSubflow = subflows.get(pathId);
        if (failedSubflow != null) {
            healthySubflows.remove(failedSubflow);
            failedSubflow.close();

            int attempts = reconnectAttempts.getOrDefault(pathId, 0);
            if (attempts < MAX_RECONNECT_ATTEMPTS) {
                reconnectAttempts.put(pathId, attempts + 1);
                scheduleReconnect(failedSubflow);
            } else {
                subflows.remove(pathId);
                lastHeartbeatTime.remove(pathId);
                reconnectAttempts.remove(pathId);
            }
        }
    }

    private void scheduleReconnect(Subflow subflow) {
        group.schedule(() -> {
            if (subflows.containsKey(subflow.getPathId())) {
                connectSubflow(subflow);
            }
        }, RECONNECT_DELAY_MS, TimeUnit.MILLISECONDS);
    }

    public void onHeartbeatReceived(int pathId) {
        AtomicLong lastTime = lastHeartbeatTime.get(pathId);
        if (lastTime != null) {
            lastTime.set(System.currentTimeMillis());
        }
    }

    public Subflow getSubflow(int pathId) {
        return subflows.get(pathId);
    }

    public void shutdown() {
        heartbeatExecutor.shutdown();
        try {
            if (!heartbeatExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                heartbeatExecutor.shutdownNow();
            }
        } catch (InterruptedException e) {
            heartbeatExecutor.shutdownNow();
            Thread.currentThread().interrupt();
        }

        subflows.values().forEach(Subflow::close);
        group.shutdownGracefully();
    }
}